package com.treasure.infrastructure.domain.llm.qianwen;

import com.alibaba.dashscope.aigc.generation.GenerationResult;
import com.alibaba.dashscope.common.ResultCallback;
import com.treasure.common.model.ChatHistoryEntity;
import com.treasure.common.model.ChatResponse;
import com.treasure.common.model.SessionUser;
import com.treasure.common.model.command.ConversationRequest;
import com.treasure.infrastructure.domain.llm.qianwen.client.QwenContext;
import com.treasure.infrastructure.model.event.ConversationEvent;
import com.treasure.common.utils.RandomIdGenerator;
import com.treasure.common.utils.TSON;
import com.treasure.domain.domain.ChatHistoryDomainService;
import com.treasure.infrastructure.client.SseEmitterServer;
import lombok.extern.slf4j.Slf4j;
import org.jetbrains.annotations.NotNull;
import org.springframework.context.ApplicationEventPublisher;
import org.springframework.context.ApplicationEventPublisherAware;

import java.time.LocalDateTime;
import java.util.Objects;
import java.util.concurrent.atomic.LongAdder;

import static com.treasure.infrastructure.domain.llm.qianwen.client.QwenGenerationClient.SSE_END_KEY;

/**
 * @author sunhongmin
 * @date 2024/1/14 10:43
 * @description
 */
@Slf4j
public class DashscopeResultCallback extends ResultCallback<GenerationResult> {
    private final QwenContext qwenContext;
    private final SessionUser sessionUser;
    private final ConversationRequest conversationRequest;
    private GenerationResult lastGenerationResult;
    private final LongAdder outputTokensCount = new LongAdder();
    private final ChatHistoryDomainService chatHistoryDomainService;
    private final ApplicationEventPublisher applicationEventPublisher;

    public DashscopeResultCallback(QwenContext qwenContext, SessionUser sessionUser, ConversationRequest conversationRequest, ChatHistoryDomainService chatHistoryDomainService, ApplicationEventPublisher applicationEventPublisher) {
        this.qwenContext = qwenContext;
        this.sessionUser = sessionUser;
        this.conversationRequest = conversationRequest;
        this.chatHistoryDomainService = chatHistoryDomainService;
        this.applicationEventPublisher = applicationEventPublisher;
    }

    @Override
    public void onEvent(GenerationResult generationResult) {
        this.lastGenerationResult = generationResult;
        String content = generationResult.getOutput().getChoices().get(0).getMessage().getContent();
        outputTokensCount.add(generationResult.getUsage().getOutputTokens());
        ChatResponse chatResponse = new ChatResponse();
        chatResponse.setInteractionId(conversationRequest.getInteractionId().getValue());
        chatResponse.setResponseMessage(content.replaceAll("\n", "<br/>"));
        SseEmitterServer.sendMessage(conversationRequest.getSseClientId().getValue(), TSON.toJsonString(chatResponse));
    }

    @Override
    public void onComplete() {
        // {"requestId":"b3aadfaf-bb9e-95e1-9e8f-05e7fd89f510","usage":{"input_tokens":869,"output_tokens":150},"output":{"choices":[{"finish_reason":"stop","message":{"role":"assistant","content":"长安大学是一所位于中国陕西省西安市的综合性大学，创建于1952年，是中国最早的林业类高等学府之一。学校占地面积约213公顷，拥有三个校区，设有八个学科门类，其中工程、植物与动物科学、农业科学和环境/生态学等学科排名在全球前1%的ESI（基本科学指标）中。长安大学也是全球前200名大学之一，根据QS世界大学排名。长安大学的国际教育学院提供各种语言课程，包括初级、中级和高级汉语课程，旨在帮助学生提高汉语水平。同时，长安大学还开设有多个留学项目，如CSC奖学金、一带一路奖学金等，欢迎广大留学生前来申请。"}}]}}
        String lastAnswer = this.lastGenerationResult.getOutput().getChoices().get(0).getMessage().getContent();
        log.info("回答： \n {}", lastAnswer);
        // 记录聊天历史
        ChatHistoryEntity chatHistoryEntity = getChatHistoryEntity(conversationRequest, lastAnswer);
        chatHistoryDomainService.addChatHistory(chatHistoryEntity);
        ChatResponse chatResponse = new ChatResponse();
        chatResponse.setInteractionId(conversationRequest.getInteractionId().getValue());
        chatResponse.setResponseMessage(SSE_END_KEY);
        SseEmitterServer.sendMessage(conversationRequest.getSseClientId().getValue(), TSON.toJsonString(chatResponse));
        // 发布拓展点 会话成功发布点，记录聊天历史、调用统计等信息
        int inputTokens = this.lastGenerationResult.getUsage().getInputTokens();
        int totalOutputTokens = outputTokensCount.intValue() - this.lastGenerationResult.getUsage().getOutputTokens();
        ConversationEvent conversationEvent = getConversationEvent(sessionUser, conversationRequest, new ConversationEvent(this), qwenContext.getQwenConfigBean().getModel(), lastAnswer, inputTokens, totalOutputTokens);
        applicationEventPublisher.publishEvent(conversationEvent);
    }

    @NotNull
    public static ConversationEvent getConversationEvent(SessionUser sessionUser, ConversationRequest conversationRequest, ConversationEvent conversationEvent, String model, String lastAnswer, int inputTokens, int totalOutputTokens) {
        if (Objects.nonNull(sessionUser)) {
            conversationEvent.setCorpId(sessionUser.getCorpId());
            conversationEvent.setUserId(sessionUser.getUserId());
        } else {
            conversationEvent.setCorpId("system");
            conversationEvent.setUserId("system");
        }
        conversationEvent.setQuestion(conversationRequest.getQuestion().getValue());
        conversationEvent.setAnswer(lastAnswer);
        conversationEvent.setCreateTime(LocalDateTime.now());
        conversationEvent.setInputTokens(inputTokens);
        conversationEvent.setOutputTokens(totalOutputTokens);
        conversationEvent.setManufacturer("DashScope");
        conversationEvent.setModel(model);
        return conversationEvent;
    }

    @Override
    public void onError(Exception e) {
        log.error("conversation error request:{},error:", TSON.toJsonString(conversationRequest), e);
        SseEmitterServer.sendMessage(conversationRequest.getSseClientId().getValue(), SSE_END_KEY);
        SseEmitterServer.removeUser(conversationRequest.getSseClientId().getValue());
    }

    @NotNull
    public static ChatHistoryEntity getChatHistoryEntity(ConversationRequest questionRequest, String result) {
        String trace = RandomIdGenerator.traceId();
        ChatHistoryEntity chatHistoryEntity = new ChatHistoryEntity();
        chatHistoryEntity.setInteractionId(questionRequest.getInteractionId());
        chatHistoryEntity.setTrace(trace);

        ChatHistoryEntity.ChatHistoryBean chatHistoryBeanForUser = new ChatHistoryEntity.ChatHistoryBean();
        chatHistoryBeanForUser.setIden(ChatHistoryEntity.IdenEnum.USER.getValue());
        chatHistoryBeanForUser.setMessage(questionRequest.getQuestion().getValue());
        chatHistoryEntity.addChatHistory(chatHistoryBeanForUser);

        ChatHistoryEntity.ChatHistoryBean chatHistoryBeanForSys = new ChatHistoryEntity.ChatHistoryBean();
        chatHistoryBeanForSys.setIden(ChatHistoryEntity.IdenEnum.ASSISTANT.getValue());
        chatHistoryBeanForSys.setMessage(result);
        chatHistoryEntity.addChatHistory(chatHistoryBeanForSys);
        return chatHistoryEntity;
    }


}
